package org.example.ycsb;

import com.alibaba.dcm.DnsCacheManipulator;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

public class HbaseYcsbSink extends RichSinkFunction<ProcessEvent> {
    private Logger logger = LoggerFactory.getLogger(HbaseYcsbSink.class);
    private transient String tableName;
    private transient int batchSize;
    private transient Connection connection;

    private transient Table table;
    private transient ConfigBean configBean;

    private transient List<ProcessEvent> buffer;

    @Override
    public void open(Configuration parameters) throws Exception {
//        updateHosts();
        ParameterTool parameterTool = ParameterTool.fromMap(getRuntimeContext().getExecutionConfig().getGlobalJobParameters().toMap());
        String ycsbConfig = parameterTool.get("configYcsbBase64");
        if (ycsbConfig == null || ycsbConfig.trim().isEmpty()) {
            configBean = new ConfigBean();
        } else {
            byte[] base64decodedBytes = Base64.getDecoder().decode(ycsbConfig);
            ycsbConfig = new String(base64decodedBytes, "utf-8");
            configBean = JackSonUtil.readValue(ycsbConfig, ConfigBean.class);
        }
        tableName = configBean.getTableName();
        batchSize = configBean.getBatchSize();
        org.apache.hadoop.conf.Configuration conf = HBaseConfiguration.create();
        conf.set("hbase.zookeeper.quorum",configBean.getHbaseUrl());
//        conf.set("hbase.zookeeper.quorum","192.168.191.14,192.168.191.15,192.168.191.16");
        conf.set("zookeeper.znode.parent","/hbase-unsecure");
        connection = ConnectionFactory.createConnection(conf);
        table = connection.getTable(TableName.valueOf(tableName));
        buffer = new ArrayList<>();
    }

    // 如果有必要,可以将HBASE集群的HOSTS写在这
    private void updateHosts(){
        String staticInfo = "192.168.xxx.xxx,xx-xx-hbase-zk01.xx.com\n" +
                "192.168.xx.xx,xx-xx-hbase-zk02.xx.com\n";
        String[] infoList = staticInfo.split("\n");
        for (String s : infoList) {
            String [] pair = s.split(",");
            String ip = pair[0];
            String hostname = pair[1];
            DnsCacheManipulator.setDnsCache(hostname, ip);
        }
    }

    @Override
    public void invoke(ProcessEvent value, Context context) {
        processElement(value,false);
    }

    private synchronized void processElement(ProcessEvent value,boolean forceFlush) {
        if(value!=null){
            buffer.add(value);
        }
        if(buffer.size()>batchSize || forceFlush){
            Exception exception = null;
            int maxRetry = 10;
            for(int i=0;i<maxRetry;i++){
                try{
                    exception=null;
                    List<KeyedTuple<Get,Put>> tuples= processBuffer(buffer);
                    flushBuffer(tuples);
                    buffer.clear();
                    break;
                }catch (Exception e){
                    logger.error("写入异常!");
                    logger.error(e.getMessage(),e);
                    exception = e;
                }
            }
            if(exception!=null){
                throw new RuntimeException(exception);
            }
        }
    }

    private void flushBuffer(List<KeyedTuple<Get,Put>> tuples) throws IOException {
        List<Get> gets = tuples.stream().map(KeyedTuple::getT1).filter(Objects::nonNull).collect(Collectors.toList());
        List<Put> puts = tuples.stream().map(KeyedTuple::getT2).filter(Objects::nonNull).collect(Collectors.toList());
        if(!gets.isEmpty()){
            table.get(gets);
        }
        if(!puts.isEmpty()){
            table.put(puts);
        }
    }

    private List<KeyedTuple<Get,Put>> processBuffer(List<ProcessEvent> buffer){
        List<KeyedTuple<Get,Put>> result = new ArrayList<>();
        for (ProcessEvent value : buffer) {
            KeyedTuple<Get,Put> tuple = new KeyedTuple<>();
            int flag = value.getFlag();
            switch (flag){
                case 1:
                    Put put = getPut(value);
                    tuple.setT2(put);
                    tuple.setKey(value.getRowKey());
                    break;
                case 2:
                    Get get = getGet(value);
                    tuple.setT1(get);
                    tuple.setKey(value.getRowKey());
                    break;
                case 3:
                    Get get1 = getGet(value);
                    Put put1 = getPut(value);
                    tuple.setT1(get1);
                    tuple.setT2(put1);
                    tuple.setKey(value.getRowKey());
                    break;
                default:
                    tuple=null;
                    break;
            }
            if(tuple!=null){
                result.add(tuple);
            }
        }
        return result;
    }

    private Put getPut(ProcessEvent value){
        String rowKey = value.getRowKey();
        if(rowKey==null || rowKey.trim().isEmpty()){
            return null;
        }
        List<String> values = value.getValues();
        Put put = new Put(Bytes.toBytes(rowKey));
        for(int i=0;i<values.size();i++){
            String columnName = String.valueOf(i);
            put.addColumn(Bytes.toBytes("cf"), Bytes.toBytes(columnName), Bytes.toBytes(values.get(i)));
        }
        return put;
    }

    private Get getGet(ProcessEvent value){
        String rowKey = value.getRowKey();
        if(rowKey==null || rowKey.trim().isEmpty()){
            return null;
        }
        return new Get(Bytes.toBytes(value.getRowKey()));
    }

    @Override
    public void close() throws Exception {
        processElement(null,true);
        table.close();
        connection.close();
    }
}
